//
// Created by yk120 on 2024/3/3.
//

#ifndef LIGHTINDEX_LIGHTWEIGHTINDEX_H
#define LIGHTINDEX_LIGHTWEIGHTINDEX_H

#include <nanofaiss/MetricType.h>

#include <cstddef>
#include <cstdint>
#include <utility>
#include <vector>

namespace faiss {

struct IndexLightWeight {
    using Entity = std::pair<float, idx_t>; // distance and id

    std::vector<uint8_t> codes;
    std::vector<idx_t> cluster; // sub -> id
    std::vector<std::vector<idx_t>> ivf; // id -> cluster id
    size_t ntotal = 0;
    size_t d;
    size_t line_size;
    bool is_trained = true;
    float sample_portion;

    IndexLightWeight(size_t d, float sample_portion = 0.01f);

    void add(size_t nx, const float *x);

    void search(const float *x, size_t nt, size_t k, float *distances, idx_t *idx, size_t probe = 1);

    // k is the size of point belong to
    void train(size_t k);
};

void random_sample(size_t n, idx_t *out, size_t sample_size, int seed = 1234);

float fvec_L2(const float *x, const float *y, size_t d);
}

#endif //LIGHTINDEX_LIGHTWEIGHTINDEX_H
